import cv2
import numpy as np
import torch


def direction_unique_group(label_slic, up_x, down_x, edge_y, tag_x, direction="x"):
    if direction == "x":
        up_index = torch.cat([up_x, edge_y], dim=1)[tag_x.squeeze(1)]
        down_index = torch.cat([down_x, edge_y], dim=1)[tag_x.squeeze(1)]
    else:
        up_index = torch.cat([edge_y, up_x], dim=1)[tag_x.squeeze(1)]
        down_index = torch.cat([edge_y, down_x], dim=1)[tag_x.squeeze(1)]
    up_label = label_slic[up_index[:, 0], up_index[:, 1]]
    down_label = label_slic[down_index[:, 0], down_index[:, 1]]
    tag = up_label != down_label
    pair_group = torch.cat([up_label[tag].unsqueeze(1), down_label[tag].unsqueeze(1)], dim=1)
    change_index = pair_group[:, 0] > pair_group[:, 1]
    change_group = torch.cat([pair_group[change_index][:, 1].unsqueeze(1), pair_group[change_index][:, 0].unsqueeze(1)],
                             dim=1)
    final_group = torch.cat([pair_group[pair_group[:, 0] < pair_group[:, 1]], change_group], dim=0).numpy()
    final_group = np.array(list(set([tuple(t) for t in final_group])))
    return final_group


def get_superpiexls_label(img):
    #初始化slic项，超像素平均尺寸20（默认为10），平滑因子20
    slic = cv2.ximgproc.createSuperpixelSLIC(img, region_size=20,ruler = 20.0)
    slic.iterate(10)     #迭代次数，越大效果越好
    mask_slic = slic.getLabelContourMask() #获取Mask，超像素边缘Mask==1
    label_slic_numpy = slic.getLabels()
    label_slic = torch.from_numpy(label_slic_numpy)       #获取超像素标签
    number_slic = slic.getNumberOfSuperpixels()  #获取超像素数目
    # 计算超像素的相邻关系
    index = torch.from_numpy(np.argwhere(mask_slic > 1))
    edge_x = index[:, 0].unsqueeze(1)
    edge_y = index[:, 1].unsqueeze(1)

    up_x = edge_x - 1
    down_x = edge_x + 1
    left_y = edge_y - 1
    right_y = edge_y + 1

    tag_up_x = (up_x >= 0) * (up_x <= img.shape[0]-1)
    tag_down_x = (down_x >= 0) * (down_x <= img.shape[0]-1)
    tag_x = tag_up_x * tag_down_x


    tag_left_y = (left_y >= 0) * (left_y <= img.shape[1]-1)
    tag_right_y = (right_y >= 0) * (right_y <= img.shape[1]-1)
    tag_y = tag_left_y * tag_right_y

    final_goup_x = direction_unique_group(label_slic, up_x, down_x, edge_y, tag_x, direction="x")

    final_goup_y = direction_unique_group(label_slic, left_y, right_y, edge_x, tag_y, direction="y")

    final_goup = np.vstack([final_goup_x, final_goup_y])
    final_group = np.array(list(set([tuple(t) for t in final_goup])))
    final_group = final_group[final_group[:, 0].argsort()]

    # mask_inv_slic = cv2.bitwise_not(mask_slic)
    # img_slic = cv2.bitwise_and(img, img, mask=mask_inv_slic)  # 在原图上绘制超像素边界
    # cv2.imshow("img_slic", img_slic)
    # cv2.waitKey(0)
    # cv2.destroyAllWindows()

    return label_slic, number_slic, final_group



